package gomoku_qtable

import (
	"bindolabs/gomoku/modules/termbox_utils"
	"encoding/gob"
	"fmt"
	"os"
	"path/filepath"
)

func (agent *QLearningAgent) LoadQTable(filename string) error {
	selectedFile := filepath.Join(SaveDir, filename)
	file, err := os.Open(selectedFile)
	if err != nil {
		return fmt.Errorf("error opening selected Q-table file: %v", err)
	}
	defer file.Close()

	decoder := gob.NewDecoder(file)
	err = decoder.Decode(&agent.qTable)
	if err != nil {
		return fmt.Errorf("error decoding Q-table: %v", err)
	}

	fmt.Printf("Successfully loaded Q-table from %s\n", selectedFile)
	return nil
}

func (agent *QLearningAgent) Load() error {
	files, err := os.ReadDir(SaveDir)
	if err != nil {
		return fmt.Errorf("error reading save directory: %v", err)
	}

	if len(files) == 0 {
		return fmt.Errorf("no saved Q-tables found in directory: %v", SaveDir)
	}

	fileNames := make([]string, 0, len(files))
	for _, file := range files {
		if !file.IsDir() {
			fileNames = append(fileNames, file.Name())
		}
	}

	selectedIndex, err := termbox_utils.DisplayMenu(fileNames)
	if err != nil {
		return err
	}

	return agent.LoadQTable(fileNames[selectedIndex])
}

func (agent *QLearningAgent) SaveQTable(filename string) error {
	file, err := os.Create(SaveDir + "/" + filename)
	if err != nil {
		return err
	}
	defer file.Close()

	encoder := gob.NewEncoder(file)
	return encoder.Encode(agent.qTable)
}
